import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
import cv2
from sklearn.preprocessing import MinMaxScaler
from Preprocessing import My_dataset
from torch.utils import data
import glob
import utils
from torchvision import datasets, transforms
import time
import torch
from Neural_Networks import (P2D_GE_Net, CNNNet, DN4Net, ProtoNetClassifier, GNNImageClassifier,
                             DSResNet_GADF, P2D_DSResNet_CWT, P2D_DSResNet_1, P2D_DSResNet_V2)
from torch import nn
import torch.nn.functional as F
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from collections import Counter
import pywt
from torch.optim.lr_scheduler import ExponentialLR
from scipy.signal import hilbert
from scipy.fft import fft
from scipy.ndimage import grey_dilation, grey_erosion, grey_opening, grey_closing
import scipy.io
from joblib import Parallel, delayed
from PIL import Image
import os
import psutil

def normalize_to_uint8(M):
    # M in [-1,1] → [0,255]
    M2 = ((M + 1) * 127.5).clip(0, 255)
    return M2.astype(np.uint8)

def GADF(signal, cataglog, count, dir):
    diff = signal[:, None] - signal[None, :]
    M = np.sin(diff)  # GADF 差场矩阵，shape = (h, h)
    GADF_image = normalize_to_uint8(M)

    save_dir = f"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\\four_data_1730\{dir}\GADF"
    os.makedirs(save_dir, exist_ok=True)

    filename = f"{count:05d}.png"
    save_path = os.path.join(save_dir, filename)

    img = Image.fromarray(GADF_image)
    img.save(save_path)
    print(filename)
    return img


def generate_cwt(signal, cataglog, count,dir, scales=np.arange(1, 65)):
    # 使用 cmor 小波
    coefficients, _ = pywt.cwt(signal, scales, 'cmor1.5-1.0')

    # 取模值并归一化到 [0,1]
    cwt_image = np.abs(coefficients)
    cwt_image = (cwt_image - np.min(cwt_image)) / (np.max(cwt_image) - np.min(cwt_image))

    # 转为 [0,255] 的 uint8 灰度图
    cwt_image_uint8 = (cwt_image * 255).astype(np.uint8)

    # 保存路径
    base_dir = f"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\\four_data_1730\{dir}\CWT"
    os.makedirs(base_dir, exist_ok=True)
    filename = f"{count:05d}.png"
    print(filename)
    save_path = os.path.join(base_dir, filename)

    # 保存为灰度图
    img = Image.fromarray(cwt_image_uint8)
    img.save(save_path)

    # 返回 numpy 而不是 PIL，方便后续处理
    return img

def function_EMTF(y, g_size):
    """单尺度形态学分解 + 权重计算"""
    structure = np.ones(g_size)

    fd = grey_dilation(y, structure=structure)
    fe = grey_erosion(y, structure=structure)
    fo = grey_opening(y, structure=structure)
    fc = grey_closing(y, structure=structure)

    # 细节波形
    WT = y - (fd + fe) / 2

    # 开闭运算组合
    fcde = grey_erosion(grey_dilation(fc, structure=structure), structure=structure)
    foed = grey_dilation(grey_erosion(fo, structure=structure), structure=structure)

    fcde_oed = (fcde + foed) / 2
    w = y - fcde_oed

    EMTF = 2 * WT * w
    return EMTF


def EMTF(signal, N, catalog, count, dir):
    """EMTF 时频图生成"""
    start_time = time.perf_counter()

    # 基础参数
    y = signal[:N]
    fs = 12000
    t = np.arange(N) / fs
    L = len(t)
    NFFT = 2 ** int(np.ceil(np.log2(L)))
    f = fs / 2 * np.linspace(0, 1, NFFT // 2 + 1)

    # 多尺度分解
    scale_range = range(1, 300)
    zatL = Parallel(n_jobs=-1)(
        delayed(function_EMTF)(y, s + 1) for s in scale_range
    )
    zatL = np.array(zatL)

    # 傅里叶包络谱
    baoluoatL = np.zeros((len(scale_range), len(f)))
    for i in range(len(scale_range)):
        analytic_signal = hilbert(zatL[i, :])
        envelope = np.abs(analytic_signal)
        envelope_fft = fft(envelope, n=NFFT) / L
        baoluoatL[i, :] = 2 * np.abs(envelope_fft[:NFFT // 2 + 1])

    # 去掉直流分量
    baoluoatL[:, 0] = 0

    # 频率范围截取 (只取前87个点)
    A = baoluoatL[:, :87]
    C_normalized = (A - np.min(A)) / (np.max(A) - np.min(A))

    # 转换成图像格式
    image_gray = (C_normalized * 255).astype(np.uint8)  # [H, W]
    image_gray = Image.fromarray(image_gray, mode='L')

    save_dir = f"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\\four_data_1730\{dir}\EMTF"
    os.makedirs(save_dir, exist_ok=True)
    filename = f"{count:05d}.png"
    save_path = os.path.join(save_dir, filename)
    image_gray.save(save_path)

    elapsed_time = time.perf_counter() - start_time
    m, s = divmod(elapsed_time, 60)
    print(f"运行第{count}次, 运行时间 {int(m)} min {s:.2f} s, "
          f"save_img: {count:03d}_{catalog}.png")

    return image_gray

def Data_loading():
    '''加载数据'''
    # 数据加载
    data_paths1 = glob.glob(r"D:\pycharm_project\Bearing_inspection\Be_legend\data\12k Drive End Bearing Fault Data\*.mat")
    data_paths2 = glob.glob(r"D:\pycharm_project\Bearing_inspection\Be_legend\data\Normal Baseline Data\*.mat")
    data_paths = data_paths1 + data_paths2

    '''数据集分类'''
    # ===============================
    # 转速：1797 rpm
    # ===============================
    Bearing_inner_ring_1797 = [105, 169, 209]  # 三个都在参考表中
    Rolling_element_1797 = [118, 185, 222]  # 3005 不在参考表，已删除
    Outer_race_1797 = [130, 144, 156, 197]  # 197,144,130,156 均存在，其他剔除
    Normal_1797 = [97]  # 保留

    # ===============================
    # 转速：1772 rpm
    # ===============================
    Bearing_inner_ring_1772 = [106, 170, 210]  # 3002 删除
    Rolling_element_1772 = [119, 186, 223]  # 3006 删除
    Outer_race_1772 = [131, 145, 158, 198]  # 保留四个，其他剔除
    Normal_1772 = [98]  # 保留

    # ===============================
    # 转速：1750 rpm
    # ===============================
    Bearing_inner_ring_1750 = [107, 171, 211]  # 全部存在，保留
    Rolling_element_1750 = [120, 187, 224]  # 3007 删除
    Outer_race_1750 = [132, 146, 159, 199]  # 保留四个（都在参考表）
    Normal_1750 = [99]  # 保留

    # ===============================
    # 转速：1730 rpm
    # ===============================
    Bearing_inner_ring_1730 = [108, 172, 212]  # 全部存在，保留
    Rolling_element_1730 = [121, 188, 225]  # 3008 删除
    Outer_race_1730 = [133, 147, 160, 200]  # 保留四个（都在参考表）
    Normal_1730 = [100]  # 保留

    data1 = {}
    data1['BIR'] = Bearing_inner_ring_1730
    data1['RE'] = Rolling_element_1730
    data1['OR'] = Outer_race_1730
    data1['ND'] = Normal_1730

    return data1, data_paths

def aLL(Bearing_inner_ring, Rolling_element, Outer_ring, nomal_data, data_paths, is_cwt=1):
    zheng = [56, 57, 58, 59, 48, 49, 50, 51]  # 应对信息处理
    count = 0
    ALL_DE_data = []
    ALL_DE_label = []
    '''数据截取'''
    # CWRU数据集中采样频率是12kHz
    fs = 12000
    # 去除前5秒的非稳态信号
    start_index = 5 * fs
    # 归一化到[-1, 1]
    scaler = MinMaxScaler(feature_range=(-1, 1))
    label_to_index = {
        'nomal_data': 0,
        'Bearing_inner_ring': 1,
        'Outer_ring': 2,
        'Rolling_element': 3,
    }

    for data_path in data_paths:
        data_split =data_path.split("\\")
        code = data_split[6].split('.')
        code_num = int(code[0])
        mat_data = sio.loadmat(data_path)

        if code_num in Bearing_inner_ring:
            if code_num in zheng:
                DE_time = mat_data[f'X0{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            else:
                DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
            ALL_DE_data.append(DE_time[start_index:])
            ALL_DE_label.append(1)

        elif code_num in Rolling_element:
            if code_num in zheng:
                DE_time = mat_data[f'X0{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            else:
                DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。

            DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
            ALL_DE_data.append(DE_time[start_index:])
            ALL_DE_label.append(3)

        elif code_num in Outer_ring:
            if code_num in zheng:
                DE_time = mat_data[f'X0{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            else:
                DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
            ALL_DE_data.append(DE_time[start_index:])
            ALL_DE_label.append(2)

        elif code_num in nomal_data:
            if code_num == 100:
                DE_time = mat_data[f'X{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            else:
                DE_time = mat_data[f'X0{code_num}_DE_time']  # 驱动端（Drive End, DE）的时域振动信号数据。
            DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()
            ALL_DE_data.append(DE_time[start_index:])
            ALL_DE_label.append(0)

    print("DE_data的有效值: " + str(len(ALL_DE_data)))
    print("DE_label的有效值: " + str(len(ALL_DE_label)))

    '''数据分段'''
    all_gadf_images = []
    all_gadf_label = []
    data_indexs = len(ALL_DE_data)
    for data_index in range(data_indexs):
        # 分段参数设置
        phi_segments_line = ALL_DE_data[data_index]
        phi_segments_line = np.clip(phi_segments_line, -1.0, 1.0)
        phi_segments_line = np.arccos(phi_segments_line)
        '''###################修改'''
        window_size = 1024      # 采样窗口
        step_size = 512        # 滑窗
        num_windows = ((phi_segments_line.shape[0] - window_size) // step_size + 1)     # 向下取整

        for num_window in range(num_windows):
            start = num_window * step_size
            end = start + window_size
            # 格拉姆角场方法（GADF）
            if is_cwt == 1:
                dir_label = f"{ALL_DE_label[data_index]}"
                M = GADF(phi_segments_line[start: end], ALL_DE_label[data_index], count, dir_label)
                # M = load_GADF_img(count)
            # CWT
            elif is_cwt == 0:
                dir_label = f"{ALL_DE_label[data_index]}"
                M = generate_cwt(phi_segments_line[start: end], ALL_DE_label[data_index], count, dir_label)
            # 形态学
            elif is_cwt == 2:
                dir_label = f"{ALL_DE_label[data_index]}"
                M = EMTF(phi_segments_line[start: end], window_size, ALL_DE_label[data_index], count, dir_label)
            else:
                raise ValueError("参数 is_cwt 必须为 0（CWT）、1（GADF）、或 2（EMTF）")

            count = count + 1
            all_gadf_images.append(M)
            all_gadf_label.append(ALL_DE_label[data_index])

    print("all_gadf_images: " + str(len(all_gadf_images)))
    print("Unique labels in training data:", np.unique(all_gadf_label))
    counts = Counter(all_gadf_label)
    print(counts[0], counts[1], counts[2], counts[3])


'''############################# 主循环 ##############################'''
if __name__ == "__main__":
    data1, data_paths = Data_loading()
    aLL(data1['BIR'], data1['RE'], data1['OR'], data1['ND'], data_paths, is_cwt=1)
    aLL(data1['BIR'], data1['RE'], data1['OR'], data1['ND'], data_paths, is_cwt=0)
    aLL(data1['BIR'], data1['RE'], data1['OR'], data1['ND'], data_paths, is_cwt=2)

